# coding=utf-8
from __future__ import absolute_import, unicode_literals

import base64
import random
import string

import six
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

from framework.utilities import get_len, to_byte, to_string


class AesCrypto(object):
    def __init__(self, encoding_aes_key):
        encoding_aes_key = AesCrypto.to_binary(encoding_aes_key + '=')
        self.cipher = AesChipher(base64.b64decode(encoding_aes_key))

    def encrypt(self, text):
        tmp_list = [to_byte(text)]
        text = b''.join(tmp_list)
        text = PKCS7Encoder.encode(text)

        ciphertext = to_byte(self.cipher.encrypt(text))
        return AesCrypto.to_text(base64.b64encode(ciphertext))

    def decrypt(self, text):
        text = to_byte(text)
        plain_text = self.cipher.decrypt(base64.b64decode(text))
        padding = plain_text[-1]
        return to_string(plain_text[:-padding])

    def _get_random_strs(self, length=16):
        rule = string.ascii_letters + string.digits
        rand_list = random.sample(rule, length)
        return ''.join(rand_list)

    @staticmethod
    def to_text(value, encoding='utf-8'):
        if not value:
            return ''
        if isinstance(value, six.text_type):
            return value
        if isinstance(value, six.binary_type):
            return value.decode(encoding)
        return six.text_type(value)

    @staticmethod
    def to_binary(value, encoding='utf-8'):
        if not value:
            return b''
        if isinstance(value, six.binary_type):
            return value
        if isinstance(value, six.text_type):
            return value.encode(encoding)
        return six.binary_type(value)


class PKCS7Encoder(object):
    block_size = 32

    @classmethod
    def encode(cls, text):
        length = get_len(text)
        padding_count = cls.block_size - length % cls.block_size
        if padding_count == 0:
            padding_count = cls.block_size
        padding = to_byte(chr(padding_count))
        return text + padding * padding_count

    @classmethod
    def decode(cls, decrypted):
        padding = decrypted[-1]
        if padding < 1 or padding > 32:
            padding = 0
        return decrypted[:-padding]


class AesChipher():
    def __init__(self, key):
        backend = default_backend()
        self.cipher = Cipher(
            algorithms.AES(key),
            modes.CBC(key[:16]),
            backend=backend
        )

    def encrypt(self, plaintext):
        encryptor = self.cipher.encryptor()
        return encryptor.update(plaintext) + encryptor.finalize()

    def decrypt(self, ciphertext):
        decryptor = self.cipher.decryptor()
        return decryptor.update(ciphertext) + decryptor.finalize()


if __name__ == '__main__':
    from framework.utilities import echo

    sEncodingAESKey = "jWmYm7qr5nMoAUwZRjGtBxmz3KA1tkAj3ykkR6q2B2s"
    echo(get_len(sEncodingAESKey), 'key')
    raw_str = to_string({"name": "董新强", "token": "dslfjoasuo;:***&$(#$@F星期五董新强哈哈哈WFDS", "age": 29})
    echo(raw_str, '原始字符串')
    aes = AesCrypto(sEncodingAESKey)
    encrypted_str = aes.encrypt(raw_str)
    echo(encrypted_str, '加密后的字符串')
    dstr = aes.decrypt(encrypted_str)
    echo(dstr, '解密后的字符串')
    echo(dstr == raw_str)
